from langchain_zhipu import ChatZhipuAI, ZhipuAIEmbeddings
from langchain_core.output_parsers import StrOutputParser
from langchain_community.vectorstores import FAISS
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.prompts import ChatPromptTemplate

from langchain.chains.combine_documents import create_stuff_documents_chain

from langchain_core.documents import Document
from langchain_community.document_loaders import WebBaseLoader
from langchain.chains import create_retrieval_chain

from langchain_community.document_loaders import TextLoader
from langchain.chains import create_history_aware_retriever
from langchain_core.prompts import MessagesPlaceholder
from langchain_core.messages import HumanMessage, AIMessage

llm = ChatZhipuAI(api_key="d3708ee404327e207b2f003775e06908.X3dgRCxbkyDfEIbh"
                  , model="glm-4")
embeddings = ZhipuAIEmbeddings(api_key="d3708ee404327e207b2f003775e06908.X3dgRCxbkyDfEIbh")

loader = WebBaseLoader("https://docs.smith.langchain.com/user_guide")

docs = loader.load()
# loader = TextLoader("index.md")
# localText = loader.load()

text_splitter = RecursiveCharacterTextSplitter()
documents = text_splitter.split_documents(docs)
print(documents)
vector = FAISS.from_documents(documents, embeddings)

retriever = vector.as_retriever()
prompt = ChatPromptTemplate.from_messages([
    MessagesPlaceholder(variable_name="chat_history"),
    ("user", "{input}"),
    ("user",
     "Given the above conversation, generate a search query to look up to get information relevant to the conversation")
])
retriever_chain = create_history_aware_retriever(llm, retriever, prompt)

# chat_history = [HumanMessage(content="Can LangSmith help test my LLM applications?"), AIMessage(content="Yes!")]
# retriever_chain.invoke({
#     "chat_history": chat_history,
#     "input": "Tell me how"
# })



prompt = ChatPromptTemplate.from_messages([
    ("system", "Answer the user's questions based on the below context:\n\n{context}"),
    MessagesPlaceholder(variable_name="chat_history"),
    ("user", "{input}"),
])
document_chain = create_stuff_documents_chain(llm, prompt)

retrieval_chain = create_retrieval_chain(retriever_chain, document_chain)
chat_history = [HumanMessage(content="Can LangSmith help test my LLM applications?"), AIMessage(content="Yes!")]
retrieval_chain.invoke({
    "chat_history": chat_history,
    "input": "Tell me how"
})